import re
import os
import time
import jieba
import random
import numpy as np
import pandas as pd
from tqdm import tqdm
from collections import Counter
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix

MAX_LEN = 200
BATCH_SIZE = 32
EMBED_DIM = 100
IS_REMOVE_STOPWORDS = False


class TextCNN(nn.Module):
    """TextCNN模型，用于文本分类任务"""

    def __init__(self, vocab_size, embed_dim, num_filters, filter_sizes, output_dim,
                 dropout=0.5, pad_idx=0, use_bn=False, static=False, pretrained_embeddings=None):
        """
        初始化 TextCNN 模型。

        Args:
            vocab_size (int): 词汇表大小。
            embed_dim (int): 词向量维度。
            num_filters (int): 每个卷积核的数量。
            filter_sizes (list): 卷积核尺寸列表。
            output_dim (int): 输出维度，即分类类别数。
            dropout (float, optional): Dropout 比率，默认为 0.5。
            pad_idx (int, optional): 填充标记的索引，默认为 0。
            use_bn (bool, optional): 是否使用批标准化，默认为 False。
            static (bool, optional): 是否冻结词向量，默认为 False。
            pretrained_embeddings (torch.Tensor, optional): 预训练词向量，默认为 None。
        """
        super().__init__()

        # 嵌入层配置，将输入的词索引转换为对应的词向量
        # vocab_size 表示词汇表大小，embed_dim 表示词向量的维度
        # padding_idx 用于指定填充标记的索引，填充位置的词向量将不会被更新
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)

        # 是否使用预训练词向量
        if pretrained_embeddings is not None:
            # 将预训练的词向量复制到嵌入层的权重中
            self.embedding.weight.data.copy_(pretrained_embeddings)
            if static:  # 是否冻结词向量
                # 若 static 为 True，则不更新词向量的权重
                self.embedding.weight.requires_grad = False

        # 卷积层配置，使用不同尺寸的卷积核提取文本特征
        # 每个卷积核的输入通道数为 1，输出通道数为 num_filters
        # 卷积核的高度为 fs，宽度为 embed_dim
        self.convs = nn.ModuleList([
            nn.Conv2d(
                in_channels=1,
                out_channels=num_filters,
                kernel_size=(fs, embed_dim)
            ) for fs in filter_sizes
        ])

        # 批标准化层
        self.use_bn = use_bn
        if use_bn:
            # 为每个卷积核的输出创建一个批标准化层
            self.bns = nn.ModuleList([nn.BatchNorm2d(num_filters) for _ in filter_sizes])

        # 全连接层和Dropout
        # 全连接层将卷积和池化后的特征映射到输出维度
        self.fc = nn.Linear(len(filter_sizes) * num_filters, output_dim)
        # Dropout 层用于防止过拟合，以指定的比率随机丢弃神经元
        self.dropout = nn.Dropout(dropout)

    def forward(self, text):
        """
        前向传播方法，定义模型的计算流程。

        Args:
            text (torch.Tensor): 输入的文本张量，形状为 [batch size, sent len]

        Returns:
            torch.Tensor: 模型的输出结果，形状为 [batch size, output_dim]
        """
        # text: [batch size, sent len]

        # 通过嵌入层将输入文本的索引转换为词向量
        embedded = self.embedding(text)
        # embedded: [batch size, sent len, emb dim]

        # 在第 1 维增加一个维度，让输入符合二维卷积层的输入要求
        embedded = embedded.unsqueeze(1)
        # embedded: [batch size, 1, sent len, emb dim]

        # 应用卷积和ReLU激活
        conved = []
        for i, conv in enumerate(self.convs):
            # 对嵌入后的文本应用卷积操作，得到特征图
            feature_map = conv(embedded)
            # 如果使用批标准化，对特征图进行批标准化处理
            if self.use_bn:
                feature_map = self.bns[i](feature_map)
            # 对特征图应用ReLU激活函数，并去除第 3 维（大小为 1）
            conved.append(F.relu(feature_map).squeeze(3))
        # conved[n]: [batch size, num_filters, sent len - filter_sizes[n] + 1]

        # 最大池化
        # 对每个卷积核的输出进行最大池化操作，提取每个特征图的最大值
        pooled = [F.max_pool1d(conv, conv.shape[2]).squeeze(2) for conv in conved]
        # pooled[n]: [batch size, num_filters]

        # 拼接不同尺寸卷积核的结果
        # 对不同卷积核池化后的结果在第 1 维上进行拼接，并应用Dropout防止过拟合
        cat = self.dropout(torch.cat(pooled, dim=1))
        # cat: [batch size, num_filters * len(filter_sizes)]

        # 通过全连接层得到最终的输出结果
        return self.fc(cat)


class TextDataset(Dataset):
    """文本数据集类，用于处理文本和标签"""
    def __init__(self, texts, labels, vocab, max_len):
        self.texts = texts
        self.labels = labels
        self.vocab = vocab
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]

        # 文本转索引
        indices = [self.vocab.get(token, self.vocab['<unk>']) for token in text]

        # 填充或截断
        if len(indices) < self.max_len:
            indices += [self.vocab['<pad>']] * (self.max_len - len(indices))
        else:
            # print(f"警告: 文本长度超过最大长度，截断: {len(indices)} > {self.max_len}")
            # print(f"文本: {text}")
            indices = indices[:self.max_len]

        return torch.tensor(indices), torch.tensor(label)

def build_vocab(texts, min_freq=2, max_size=None):
    """构建词汇表

    Args:
        texts: 分词后的文本列表
        min_freq: 最小词频，低于此频率的词将被忽略
        max_size: 词汇表最大大小，None表示不限制

    Returns:
        词汇表字典，键为词，值为索引
    """
    vocab = {'<pad>': 0, '<unk>': 1}
    word_freq = Counter()

    for text in texts:
        word_freq.update(text)

    # 按词频排序
    sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)

    # 构建词汇表
    for word, freq in sorted_words:
        if freq >= min_freq and (max_size is None or len(vocab) < max_size):
            vocab[word] = len(vocab)

    return vocab

def preprocess_text(texts, language='zh', remove_stopwords=True, stopwords_path=None, max_len=MAX_LEN):
    """文本预处理

    Args:
        texts: 原始文本列表
        language: 语言类型，'zh'表示中文，其他表示英文
        remove_stopwords: 是否移除停用词
        stopwords_path: 停用词文件路径

    Returns:
        分词后的文本列表
    """
    processed_texts = []

    # 加载停用词
    stopwords = set()
    if remove_stopwords and stopwords_path:
        try:
            with open(stopwords_path, 'r', encoding='utf-8') as f:
                stopwords = set([line.strip() for line in f])
        except Exception as e:
            print(f"警告: 加载停用词文件失败: {e}")
            stopwords = set()

    for text in texts:
        # 移除特殊字符
        if language == 'zh':
            text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9]', ' ', text)
        else:
            text = re.sub(r'[^\w\s]', '', text)

        # 分词
        if language == 'zh':  # 中文文本
            # tokens = jieba.lcut(text.lower())
            # 无效果
            # tokens = jieba.lcut_for_search(text.lower())
            tokens = list(text.lower())
        else:  # 英文文本
            tokens = text.lower().split()

        # 移除停用词
        if remove_stopwords:
            tokens = [token for token in tokens if token not in stopwords]

        if len(tokens) > max_len:
            print(f"警告: 文本长度超过最大长度，截断: {len(tokens)} > {max_len}")
            print(f"文本: {text}")

        processed_texts.append(tokens)

    return processed_texts

def load_data_from_csv(file_path, text_col, label_col, language='zh',
                       remove_stopwords=True, stopwords_path=None,
                       sample_size=None, random_state=42, max_len=MAX_LEN):
    """从CSV文件加载数据

    Args:
        file_path: CSV文件路径
        text_col: 文本列名
        label_col: 标签列名
        language: 语言类型
        remove_stopwords: 是否移除停用词
        stopwords_path: 停用词文件路径
        sample_size: 采样大小，None表示使用全部数据
        random_state: 随机种子

    Returns:
        处理后的文本列表，标签列表
    """
    # 读取CSV文件
    try:
        df = pd.read_csv(file_path)
        print(f"成功加载数据: {len(df)} 条记录")
    except Exception as e:
        print(f"加载数据失败: {e}")
        return [], []

    # 采样数据
    if sample_size and sample_size < len(df):
        df = df.sample(sample_size, random_state=random_state)
        print(f"采样数据: {len(df)} 条记录")

    # 提取文本和标签
    texts = df[text_col].astype(str).tolist()
    labels = df[label_col].tolist()

    # 预处理文本
    processed_texts = preprocess_text(
        texts,
        language=language,
        remove_stopwords=remove_stopwords,
        stopwords_path=stopwords_path,
        max_len=max_len,
    )

    return processed_texts, labels

def train_model(model, train_loader, optimizer, criterion, device, scheduler=None):
    """训练模型一个 epoch

    Args:
        model: 模型实例
        train_loader: 训练数据加载器
        optimizer: 优化器
        criterion: 损失函数
        device: 计算设备
        scheduler: 学习率调度器

    Returns:
        平均训练损失
    """
    # 将模型设置为训练模式，启用 Dropout 等训练时才使用的机制
    model.train()
    total_loss = 0

    for texts, labels in tqdm(train_loader, desc='训练中'):
        texts, labels = texts.to(device), labels.to(device)

        # 清空优化器中存储的梯度信息，防止梯度累积
        optimizer.zero_grad()
        # 前向传播，将文本数据输入模型得到预测结果
        predictions = model(texts)
        # 计算预测结果与真实标签之间的损失
        loss = criterion(predictions, labels)
        # 反向传播，计算损失函数关于模型参数的梯度
        loss.backward()

        # 梯度裁剪，防止梯度爆炸，将模型参数的梯度范数限制在最大为 1.0 的范围内
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # 根据计算得到的梯度更新模型参数
        optimizer.step()
        # 如果提供了学习率调度器，则更新学习率
        if scheduler:
            scheduler.step()

        total_loss += loss.item()

    return total_loss / len(train_loader)

def evaluate_model(model, data_loader, criterion, device):
    """评估模型性能

    Args:
        model: 模型实例
        data_loader: 数据加载器
        criterion: 损失函数
        device: 计算设备

    Returns:
        平均损失，预测结果列表，真实标签列表
    """
    model.eval()
    total_loss = 0
    all_predictions = []
    all_labels = []

    with torch.no_grad():
        for texts, labels in tqdm(data_loader, desc='评估中'):
            texts, labels = texts.to(device), labels.to(device)

            predictions = model(texts)
            loss = criterion(predictions, labels)

            total_loss += loss.item()
            all_predictions.extend(predictions.argmax(dim=1).cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    return total_loss / len(data_loader), all_predictions, all_labels

def predict_text(model, text, vocab, max_len, language='zh',
                 remove_stopwords=True, stopwords_path=None, device='cpu'):
    """对单个文本进行分类预测

    Args:
        model: 训练好的模型
        text: 待预测的文本
        vocab: 词汇表
        max_len: 最大文本长度
        language: 语言类型
        remove_stopwords: 是否移除停用词
        stopwords_path: 停用词文件路径
        device: 计算设备

    Returns:
        预测类别，预测概率
    """
    model.eval()

    # 预处理文本
    processed_text = preprocess_text(
        [text],
        language=language,
        remove_stopwords=remove_stopwords,
        stopwords_path=stopwords_path
    )[0]

    # 文本转索引
    indices = [vocab.get(token, vocab['<unk>']) for token in processed_text]

    # 填充或截断
    if len(indices) < max_len:
        indices += [vocab['<pad>']] * (max_len - len(indices))
    else:
        start = random.choice(range(len(indices) - max_len + 1))
        indices = indices[start:start+max_len]

    # 转换为张量
    input_tensor = torch.tensor([indices]).to(device)

    # 预测
    with torch.no_grad():
        output = model(input_tensor)
        probs = F.softmax(output, dim=1)
        pred = torch.argmax(probs, dim=1).item()
        prob = probs[0][pred].item() # type: ignore

    return pred, prob

def run_textcnn(data_file, text_col='text', label_col='label', test_size=0.2,
                embed_dim=100, num_filters=100, filter_sizes=[3, 4, 5],
                output_dim=2, dropout=0.5, batch_size=64, epochs=10,
                max_len=100, language='zh', use_bn=False, static=False,
                lr=1e-3, weight_decay=0., early_stopping=3,
                remove_stopwords=True, stopwords_path=None,
                sample_size=None, random_state=42, save_model_path='best_model.pt',
                pretrained_embeddings=None):
    """运行TextCNN模型训练和评估流程

    Args:
        data_file: 数据CSV文件路径
        text_col: 文本列名
        label_col: 标签列名
        test_size: 测试集比例
        embed_dim: 词向量维度
        num_filters: 每个卷积核的数量
        filter_sizes: 卷积核尺寸列表
        output_dim: 输出维度，即分类类别数
        dropout: Dropout比率
        batch_size: 批次大小
        epochs: 训练轮数
        max_len: 文本最大长度
        language: 语言类型，'zh'表示中文，其他表示英文
        use_bn: 是否使用批标准化
        static: 是否冻结词向量
        lr: 学习率
        weight_decay: L2正则化系数
        early_stopping: 早停轮数，连续多少轮性能不提升则停止训练
        remove_stopwords: 是否移除停用词
        stopwords_path: 停用词文件路径
        sample_size: 数据采样大小
        random_state: 随机种子
        save_model_path: 保存模型的路径
        pretrained_embeddings: 预训练词向量

    Returns:
        训练好的模型，词汇表
    """
    # 设置随机种子，保证结果可复现
    random.seed(random_state)
    np.random.seed(random_state)
    torch.manual_seed(random_state)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(random_state)

    # 加载数据
    print("正在加载数据...")
    texts, labels = load_data_from_csv(
        data_file,
        text_col,
        label_col,
        language=language,
        remove_stopwords=remove_stopwords,
        stopwords_path=stopwords_path,
        sample_size=sample_size,
        random_state=random_state,
        max_len=max_len,
    )

    if not texts:
        print("没有有效的数据，程序退出")
        return None, None

    # 划分训练集和测试集
    print(f"正在将数据划分为训练集({1-test_size})和测试集({test_size})...")
    train_texts, test_texts, train_labels, test_labels = train_test_split(
        texts, labels, test_size=test_size, random_state=random_state,
        stratify=labels if output_dim > 1 else None
    )

    test_len = len(test_texts)
    print(f"训练集大小: {len(train_texts)}")
    print(f"测试集大小: {test_len}")

    # 构建词汇表
    print("正在构建词汇表...")
    vocab = build_vocab(train_texts)
    vocab_size = len(vocab)
    print(f"词汇表大小: {vocab_size}")

    # 划分训练集和验证集
    train_texts_proc, val_texts_proc, train_labels, val_labels = train_test_split(
        train_texts, train_labels, test_size=0.1, random_state=random_state,
        stratify=train_labels if output_dim > 1 else None
    )

    # 创建数据集和数据加载器
    train_dataset = TextDataset(train_texts_proc, train_labels, vocab, max_len)
    val_dataset = TextDataset(val_texts_proc, val_labels, vocab, max_len)
    test_dataset = TextDataset(test_texts, test_labels, vocab, max_len)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)

    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # 初始化模型
    model = TextCNN(
        vocab_size=vocab_size,
        embed_dim=embed_dim,
        num_filters=num_filters,
        filter_sizes=filter_sizes,
        output_dim=output_dim,
        dropout=dropout,
        use_bn=use_bn,
        static=static,
        pretrained_embeddings=pretrained_embeddings
    ).to(device)

    # 定义优化器和损失函数
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    criterion = nn.CrossEntropyLoss()

    # 学习率调度器和早停机制
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=1,  #verbose=True
    )
    best_train_loss = float('inf')
    best_val_loss = float('inf')
    early_stop_count = 0

    # 训练模型
    print("开始训练模型...")
    start_time = time.time()
    for epoch in range(epochs):
        train_loss = train_model(model, train_loader, optimizer, criterion, device)
        val_loss, _, _ = evaluate_model(model, val_loader, criterion, device)

        # 更新学习率
        scheduler.step(val_loss)

        print(f'轮次: {epoch+1}/{epochs}')
        print(f'训练损失: {train_loss:.4f}')
        print(f'验证损失: {val_loss:.4f}')
        print('-' * 50)

        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_train_loss = train_loss
            torch.save(model.state_dict(), save_model_path)
            print(f"模型已保存到 {save_model_path}")
            early_stop_count = 0
        else:
            early_stop_count += 1
            if early_stop_count >= early_stopping:
                print(f"早停: 在第 {epoch+1} 轮停止训练, 最佳损失: {best_train_loss:.4f}, {best_val_loss:.4f}")
                break

    # 加载最佳模型并在测试集上评估
    model.load_state_dict(torch.load(save_model_path))
    start_time = time.time()
    test_loss, predictions, labels = evaluate_model(model, test_loader, criterion, device)
    end_time = time.time()

    print("预测错误的记录:")
    for i, (pred, label) in enumerate(zip(predictions, labels)):
        if pred != label:
            print(f"索引: {i}, 预测: {pred}, 真实: {label}, 长度: {len(test_texts[i])}, 文本: {"".join(test_texts[i])}")

    print(f"测试耗时: {end_time - start_time:.3f} 秒, 单条记录预测耗时: {(end_time - start_time) * 1000 / test_len:.2f} 毫秒")
    print(f'训练耗时: {end_time - start_time:.0f} 秒')
    print(f'测试损失: {test_loss:.4f}')
    print(f'测试准确率: {(np.array(predictions) == np.array(labels)).mean():.3f}')

    # 打印分类报告
    print("分类报告:")
    print(classification_report(labels, predictions))

    # 打印混淆矩阵
    print("混淆矩阵:")
    print(confusion_matrix(labels, predictions))

    return model, vocab

# 使用示例 - 从CSV文件加载数据并拆分训练集和测试集
if __name__ == "__main__":
    import sys
    # 示例数据文件路径（请替换为实际文件路径）
    data_file = sys.argv[1]

    # 创建示例数据（实际使用时请替换为真实数据）
    if not os.path.exists(data_file):
        raise  Exception("示例数据文件不存在，请替换为实际数据文件路径")

    # 运行模型
    model, vocab = run_textcnn(
        data_file=data_file,
        text_col='review',
        label_col='label',
        test_size=0.2,  # 将数据按8:2划分为训练集和测试集
        embed_dim=EMBED_DIM,
        num_filters=100,
        filter_sizes=[3, 4, 5],
        output_dim=2,
        dropout=0.5,
        batch_size=BATCH_SIZE,
        epochs=20,
        max_len=MAX_LEN,
        language='zh',
        use_bn=True,
        static=False,
        lr=1e-3,
        weight_decay=1e-4,
        early_stopping=3,
        remove_stopwords=IS_REMOVE_STOPWORDS,
        random_state=42,
        save_model_path='textcnn_best_model.pt'
    )

    # 示例：对新文本进行预测
    if model and vocab:
        test_text = "这是一部非常精彩的电影，值得一看"
        pred, prob = predict_text(
            model,
            test_text,
            vocab,
            max_len=MAX_LEN,
            language='zh',
            remove_stopwords=True,
            device='cuda' if torch.cuda.is_available() else 'cpu'
        )
        print(f'预测结果: 类别 {pred}，置信度 {prob:.4f}')
